# -*- coding:utf-8 -*-

# @Time    : 2023/3/23 21:39
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : merge.py

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
import re
import json
import tqdm
import utils
import random
#读取all.json
def read_json(path):
    try:
        data = utils.jload(path)
        new_data = []
        for instruct_dict in tqdm.tqdm(data):
            # 只保留data的instruction, input, output
            instruct_dict_new = {}
            instruct_dict_new["instruction"] = instruct_dict.get("instruction", "")
            instruct_dict_new["input"] = instruct_dict.get("input", "")
            instruct_dict_new["output"] = instruct_dict.get("output","")
            new_data.append(instruct_dict_new)

    except:
        # 按行读取，每行是个json
        '''with open(path, "r", encoding="utf-8") as f:
            new_data = list()
            for line in f:
                try:
                    instruct_dict = json.loads(line)
                    input = instruct_dict["input"]
                    target = instruct_dict["target"]
                    # 对于文本中包含bell的数据删掉整条数据，大小写不敏感
                    if "bell" in input.lower() or "bell" in target.lower():
                        continue
                    instruct_dict_new = {}
                    instruct_dict_new["instruction"] = input
                    instruct_dict_new["output"] = target

                    new_data.append(instruct_dict_new)
                except:
                    print("red_json error:")
                    print(line)
                    continue'''

        # 读取BELLE生成数据
        new_data = list()
        with open(path, 'r') as fp:
            lines = fp.readlines()
            for line in tqdm.tqdm(lines):
                # data.append(json.loads(line[:-1]))

                if len(json.loads(line[:-1])['input'].split('\n')) == 2 and \
                        "belle" not in json.loads(line[:-1])['input'].split('\n')[0].lower() and \
                        "belle" not in json.loads(line[:-1])['input'].split('\n')[1].lower():

                    new_data.append({
                        'instruction': json.loads(line[:-1])['input'].split('\n')[0],
                        'input': json.loads(line[:-1])['input'].split('\n')[1],
                        'output': json.loads(line[:-1])['target']
                    })

    return new_data

#写入json,用utf-8编码
def write_json(path, data):
    utils.jdump(data, path)


#合并两个json文件
def merge_json(file_paths, output_path):

    all_data = list()
    for file_path in file_paths:
        data = read_json(file_path)
        all_data += data

    # 随机打乱数据
    random.shuffle(all_data)
    write_json(output_path, all_data)


# 转换结果成chatllama的格式instruction对应user_input，output对应completion
def convert_format(path1, path2):
    data = read_json(path1)
    new_data = []
    for instruct_dict in data:
        instruct_dict_new = {}
        instruct_dict_new["user_input"] = instruct_dict.get("instruction", "")
        if "input" in instruct_dict:
            instruct_dict_new["user_input"] = instruct_dict_new["user_input"]  + "\n" + instruct_dict.get("input", "")
        instruct_dict_new["completion"] = instruct_dict.get("output", "")
        new_data.append(instruct_dict_new)
    write_json(path2, new_data)


# BELLE 生成数据格式转换，清洗
def transfer_belle_data():

    file_paths = ['./data/Belle.train.json', './data/belle_open_source_1M.train.json']

    new_data = list()
    for file_path in tqdm.tqdm(file_paths):
        with open(file_path, 'r') as fp:
            lines = fp.readlines()
            for line in tqdm.tqdm(lines):
                # data.append(json.loads(line[:-1]))

                if len(json.loads(line[:-1])['input'].split('\n')) == 2 and \
                    "belle" not in json.loads(line[:-1])['input'].split('\n')[0].lower() and \
                    "belle" not in json.loads(line[:-1])['input'].split('\n')[1].lower() and \
                    "法语" not in json.loads(line[:-1])['input'].split('\n')[0].lower() and \
                    "法文" not in json.loads(line[:-1])['input'].split('\n')[0].lower() :
                    # print(json.loads(line[:-1])['input'].split('\n'))

                    new_data.append({
                        'instruction': json.loads(line[:-1])['input'].split('\n')[0],
                        'input': json.loads(line[:-1])['input'].split('\n')[1],
                        'output': json.loads(line[:-1])['target']
                    })

    with open('./data/all_belle.train_transfer.json', 'w', encoding='utf-8') as fp:
        utils.jdump(new_data, fp)


# 获取文件内的数据量
def get_num(file_path):

    machine_instruction_data = utils.jload(file_path)
    print( len(machine_instruction_data) )


def task():
    file_paths = list()
    for i in range(1, 11):
        file_paths.append("./output/jinrong_multi/lingxi.train.jinrong_{}.json".format(i))
        file_paths.append("./output/baoxian_multi/lingxi.train.baoxian_{}.json".format(i))

    # file_paths.extend(['./data/lingxi_intent_single.json', './data/lingxi_intent_multi.json',  './data/lingxi_intent_meaning.json'])
    file_paths.extend(['./data/alpaca_data.json', './data/all_belle.train_transfer.json'])
    merge_json(file_paths,
                './output/result_merge.json')

    get_num('./output/result_merge.json')

def encode_prompt(prompt_instructions):
    """Encode multiple prompt instructions into a single string."""
    prompt = open("./seed/prompt_cn.txt").read() + "\n"

    for idx, task_dict in enumerate(prompt_instructions):
        (instruction, input, output) = task_dict["instruction"], task_dict["input"], task_dict["output"]
        instruction = re.sub(r"\s+", " ", instruction).strip().rstrip(":")
        input = "<无输入>" if input.lower() == "" else input
        prompt += f"###\n"
        prompt += f"{idx + 1}. 指令: {instruction}\n"
        prompt += f"{idx + 1}. 输入:\n{input}\n"
        prompt += f"{idx + 1}. 输出:\n{output}\n"
    prompt += f"###\n"
    prompt += f"{idx + 2}. 指令:"
    return prompt


if __name__ == "__main__":


    # task()

    '''file_paths = list()
    for i in range(1, 11):
        file_paths.append("./output/insurance_planner_multi/lingxi.train.insurance_planner_{}.json".format(i))

    merge_json(file_paths, './output/insurance_planner_result_merge.json')'''



    '''merge_json(
        ['./dialogue_data/2186_0509.json', './dialogue_data/2186_0509_intent.json', './dialogue_data/2186_0509_no_intent.json',
        './dialogue_data/2199_0509.json', './dialogue_data/2199_0509_intent.json', './dialogue_data/2199_0509_no_intent.json', ],
        'train_0509_dialogue.json')'''
    merge_json(['./dialogue_data/2226_0509_1.json','./dialogue_data/2226_0509_intent.json', './dialogue_data/2226_0509_no_intent.json',],
               './dialogue_data/2226_0509_all.json')

    # merge_json(
    #     ['./dialogue_data/2186_0509_1.json', './dialogue_data/2186_0509_2.json'], './dialogue_data/2186_0509.json',
    # )

